# 接收端代码 - 使用CircuitPython
# 功能：接收WiFi控制信号，驱动直流电机

import board
import wifi
import socketpool
import time
import pwmio
import digitalio

# WiFi热点配置
AP_SSID = "esp32_motor_ap"
AP_PASSWORD = "12345678"
UDP_PORT = 4210

# 选择驱动方案（取消注释选择一种）
DRIVER_TYPE = "DUAL_MOS"  # 可选: "H_BRIDGE", "RELAY", "DUAL_MOS", "SINGLE_MOS"

# 全局变量声明
motor_a1 = None
motor_a2 = None
motor_b1 = None
motor_b2 = None
motor_pwm = None
motor_pwm1 = None
motor_pwm2 = None
direction_relay = None

def initialize_driver():
    """根据选择的驱动方案初始化引脚"""
    global motor_a1, motor_a2, motor_b1, motor_b2
    global motor_pwm, motor_pwm1, motor_pwm2, direction_relay
    
    if DRIVER_TYPE == "H_BRIDGE":
        # H桥驱动 - 4个引脚
        motor_a1 = digitalio.DigitalInOut(board.IO16)
        motor_a2 = digitalio.DigitalInOut(board.IO17)
        motor_b1 = digitalio.DigitalInOut(board.IO18)
        motor_b2 = digitalio.DigitalInOut(board.IO19)
        
        for pin in [motor_a1, motor_a2, motor_b1, motor_b2]:
            pin.direction = digitalio.Direction.OUTPUT
            pin.value = False
        print("H桥驱动初始化完成")
        
    elif DRIVER_TYPE == "RELAY":
        # 继电器方案 - 1个PWM + 1个方向控制
        motor_pwm = pwmio.PWMOut(board.IO18, frequency=20000, duty_cycle=0)
        direction_relay = digitalio.DigitalInOut(board.IO19)
        direction_relay.direction = digitalio.Direction.OUTPUT
        print("继电器驱动初始化完成")
        
    elif DRIVER_TYPE == "DUAL_MOS":
        # 双MOS管方案 - 2个PWM
        motor_pwm1 = pwmio.PWMOut(board.IO18, frequency=20000, duty_cycle=0)
        motor_pwm2 = pwmio.PWMOut(board.IO19, frequency=20000, duty_cycle=0)
        print("双MOS管驱动初始化完成")
        
    elif DRIVER_TYPE == "SINGLE_MOS":
        # 单MOS管方案 - 只能控制转速，无法改变方向
        motor_pwm = pwmio.PWMOut(board.IO18, frequency=20000, duty_cycle=0)
        print("警告：单MOS管无法控制方向，将忽略方向信号")
        print("单MOS管驱动初始化完成")
    else:
        print(f"错误：未知的驱动方案 {DRIVER_TYPE}")
        return False
    
    return True

def control_motor(direction, speed):
    """根据选择的驱动方案控制电机"""
    if DRIVER_TYPE == "H_BRIDGE":
        control_motor_h_bridge(direction, speed)
    elif DRIVER_TYPE == "RELAY":
        control_motor_relay(direction, speed)
    elif DRIVER_TYPE == "DUAL_MOS":
        control_motor_dual_mos(direction, speed)
    elif DRIVER_TYPE == "SINGLE_MOS":
        control_motor_single_mos(speed)  # 忽略方向
    else:
        print(f"错误：未初始化的驱动方案 {DRIVER_TYPE}")

def control_motor_h_bridge(direction, speed):
    """H桥控制方案"""
    global motor_a1, motor_a2, motor_b1, motor_b2
    
    if not all([motor_a1, motor_a2, motor_b1, motor_b2]):
        print("错误：H桥驱动未正确初始化")
        return
    
    # 停止所有引脚
    motor_a1.value = False
    motor_a2.value = False
    motor_b1.value = False
    motor_b2.value = False
    
    if direction == 1:  # 正转
        motor_a1.value = True
        motor_b1.value = True
        print(f"电机正转，速度: {speed}")
    elif direction == -1:  # 反转
        motor_a2.value = True
        motor_b2.value = True
        print(f"电机反转，速度: {speed}")
    else:  # 停止
        print("电机停止")

def control_motor_relay(direction, speed):
    """继电器控制方案"""
    global motor_pwm, direction_relay
    
    if not motor_pwm or not direction_relay:
        print("错误：继电器驱动未正确初始化")
        return
    
    if direction == 1:  # 正转
        direction_relay.value = True
        duty_cycle = int(speed * 65535 / 255)
        motor_pwm.duty_cycle = duty_cycle
        print(f"电机正转，速度: {speed}")
    elif direction == -1:  # 反转
        direction_relay.value = False
        duty_cycle = int(speed * 65535 / 255)
        motor_pwm.duty_cycle = duty_cycle
        print(f"电机反转，速度: {speed}")
    else:  # 停止
        motor_pwm.duty_cycle = 0
        print("电机停止")

def control_motor_dual_mos(direction, speed):
    """双MOS管控制方案"""
    global motor_pwm1, motor_pwm2
    
    if not motor_pwm1 or not motor_pwm2:
        print("错误：双MOS管驱动未正确初始化")
        return
    
    motor_pwm1.duty_cycle = 0
    motor_pwm2.duty_cycle = 0
    
    if direction == 1:  # 正转
        duty_cycle = int(speed * 65535 / 255)
        motor_pwm1.duty_cycle = duty_cycle
        print(f"电机正转，速度: {speed}")
    elif direction == -1:  # 反转
        duty_cycle = int(speed * 65535 / 255)
        motor_pwm2.duty_cycle = duty_cycle
        print(f"电机反转，速度: {speed}")
    else:  # 停止
        print("电机停止")

def control_motor_single_mos(speed):
    """单MOS管控制方案（只能控制转速）"""
    global motor_pwm
    
    if not motor_pwm:
        print("错误：单MOS管驱动未正确初始化")
        return
    
    if speed > 0:
        duty_cycle = int(speed * 65535 / 255)
        motor_pwm.duty_cycle = duty_cycle
        print(f"电机运行，速度: {speed}")
    else:
        motor_pwm.duty_cycle = 0
        print("电机停止")

def setup_access_point():
    """设置WiFi热点"""
    print("正在启动WiFi热点...")
    
    try:
        # 配置AP模式
        wifi.radio.start_ap(AP_SSID, AP_PASSWORD)
        
        print(f"WiFi热点已启动")
        print(f"SSID: {AP_SSID}")
        print(f"密码: {AP_PASSWORD}")
        print(f"IP地址: {wifi.radio.ipv4_address}")
        return True
    except Exception as e:
        print(f"WiFi热点启动失败: {e}")
        return False

def parse_control_data(data):
    """解析控制数据"""
    try:
        # 数据格式: "方向,速度"
        parts = data.split(',')
        if len(parts) == 2:
            direction = int(parts[0])
            speed = int(parts[1])
            
            # 验证数据范围
            if direction not in [-1, 0, 1]:
                print(f"方向值无效: {direction}")
                return 0, 0
            if speed < 0 or speed > 255:
                print(f"速度值无效: {speed}")
                return 0, 0
                
            return direction, speed
        else:
            print(f"数据格式错误: {data}")
            return 0, 0
    except ValueError as e:
        print(f"数据解析错误: {e}")
        return 0, 0

def receive_control_data():
    """接收控制数据"""
    try:
        # 创建UDP socket
        pool = socketpool.SocketPool(wifi.radio)
        sock = pool.socket()
        sock.bind(("0.0.0.0", UDP_PORT))
        
        print(f"开始监听UDP端口 {UDP_PORT}...")
        
        while True:
            try:
                # 接收数据
                data, addr = sock.recvfrom(1024)
                if data:
                    # 解析数据
                    control_data = data.decode().strip()
                    direction, speed = parse_control_data(control_data)
                    
                    # 控制电机
                    control_motor(direction, speed)
                    
                    print(f"收到来自 {addr} 的数据: {control_data}")
                    
            except Exception as e:
                print(f"接收数据错误: {e}")
                time.sleep(0.1)
                
    except Exception as e:
        print(f"Socket创建失败: {e}")
    finally:
        if 'sock' in locals():
            sock.close()

def stop_all_motors():
    """停止所有电机"""
    print("停止所有电机...")
    
    if DRIVER_TYPE == "H_BRIDGE":
        if all([motor_a1, motor_a2, motor_b1, motor_b2]):
            motor_a1.value = False
            motor_a2.value = False
            motor_b1.value = False
            motor_b2.value = False
    elif DRIVER_TYPE == "RELAY":
        if motor_pwm:
            motor_pwm.duty_cycle = 0
    elif DRIVER_TYPE == "DUAL_MOS":
        if motor_pwm1 and motor_pwm2:
            motor_pwm1.duty_cycle = 0
            motor_pwm2.duty_cycle = 0
    elif DRIVER_TYPE == "SINGLE_MOS":
        if motor_pwm:
            motor_pwm.duty_cycle = 0

def main():
    """主程序"""
    print("ESP32电机控制器启动...")
    print(f"驱动方案: {DRIVER_TYPE}")
    
    # 初始化驱动
    if not initialize_driver():
        print("驱动初始化失败，程序退出")
        return
    
    # 设置WiFi热点
    if not setup_access_point():
        print("WiFi热点设置失败，程序退出")
        return
    
    # 初始化电机控制
    print("初始化电机控制...")
    stop_all_motors()  # 确保电机停止
    
    try:
        # 开始接收控制数据
        receive_control_data()
    except KeyboardInterrupt:
        print("程序被用户中断")
    except Exception as e:
        print(f"程序运行错误: {e}")
    finally:
        # 程序结束时停止所有电机
        stop_all_motors()
        print("程序结束，所有电机已停止")

if __name__ == "__main__":
    main()